#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
文档解析器主程序
支持单文件和批量文件处理模式
"""

import os
import sys
from typing import List, Tuple, Optional
from docs_parser import parse_document, save_to_json


def get_document_files(docs_dir: str) -> Tuple[List[str], List[str], List[str]]:
    """获取文档目录中的所有支持的文档文件"""
    if not os.path.exists(docs_dir):
        print(f"错误: {docs_dir}目录不存在")
        return [], [], []

    supported_extensions = ['.docx', '.pdf']
    all_files = []
    docx_files = []
    pdf_files = []

    for file in os.listdir(docs_dir):
        if any(file.lower().endswith(ext) for ext in supported_extensions):
            all_files.append(file)
            if file.lower().endswith('.docx'):
                docx_files.append(file)
            elif file.lower().endswith('.pdf'):
                pdf_files.append(file)

    return all_files, docx_files, pdf_files


def display_document_list(all_files: List[str]) -> bool:
    """显示文档列表"""
    if not all_files:
        print("错误: docs目录中没有找到支持的文档文件(.docx或.pdf)")
        return False

    print(f"\n找到 {len(all_files)} 个文档文件:")
    for i, file in enumerate(all_files, 1):
        print(f"  {i}. {file}")
    return True


def get_user_choice(all_files: List[str]) -> Optional[str]:
    """获取用户选择的文件"""
    if len(all_files) > 1:
        try:
            choice = int(input(f"\n请选择要解析的文档编号 (1-{len(all_files)}): "))
            if choice < 1 or choice > len(all_files):
                print("无效的选择")
                return None
            return all_files[choice - 1]
        except ValueError:
            print("请输入有效的数字")
            return None
    else:
        selected_file = all_files[0]
        print(f"\n自动选择文档: {selected_file}")
        return selected_file


def process_document(file_path: str, output_dir: str = "result"):
    """处理单个文档"""
    file_name = os.path.basename(file_path)
    file_ext = os.path.splitext(file_name)[1].lower()

    print(f"\n正在处理文档: {file_name}")

    try:
        # 根据文件后缀调用相应的parser
        data = parse_document(file_path)

        # 保存JSON到result目录
        os.makedirs(output_dir, exist_ok=True)
        base_name = os.path.splitext(file_name)[0]
        output_filename = f"{base_name}.json"
        output_path = os.path.join(output_dir, output_filename)

        save_to_json(data, output_path)

        print("\n解析完成!")
        print(f"共找到 {len(data['sections'])} 个主要章节")
        print(f"结果已保存到: {output_path}")

    except Exception as e:
        print(f"处理文档时出错: {e}")


def batch_mode_main():
    """批量处理模式主函数"""
    docs_dir = "docs"
    
    # 获取所有支持的文件
    all_files, docx_files, pdf_files = get_document_files(docs_dir)
    if not all_files:
        print("错误: docs目录中没有找到支持的文档文件")
        return
        
    # 动态识别支持的文件后缀
    supported_extensions = set()
    for file in all_files:
        ext = os.path.splitext(file)[1].lower()
        supported_extensions.add(ext)
        
    print("\n文档解析器 - 批量文件模式")
    print("=" * 50)
    print(f"支持的文件后缀: {', '.join(supported_extensions)}")
    
    # 选择处理方式
    print("\n处理选项:")
    print("1. 处理所有支持的文件")
    print("2. 只处理 .docx 文件")
    print("3. 只处理 .pdf 文件")
    print("4. 返回主菜单")
    
    choice = input("请选择处理方式 (1-4): ").strip()
    
    if choice == "1":
        files_to_process = all_files
    elif choice == "2":
        files_to_process = docx_files if docx_files else []
    elif choice == "3":
        files_to_process = pdf_files if pdf_files else []
    elif choice == "4":
        return
    else:
        print("无效选择")
        return
        
    if not files_to_process:
        print("没有找到符合条件的文件")
        return
        
    # 处理选中的文件
    result_dir = "result"
    os.makedirs(result_dir, exist_ok=True)
    
    success_count = 0
    for i, file_name in enumerate(files_to_process, 1):
        print(f"\n[{i}/{len(files_to_process)}] 处理文档: {file_name}")
        try:
            file_path = os.path.join(docs_dir, file_name)
            
            # 根据文件后缀调用相应的parser
            data = parse_document(file_path)
            
            # 保存结果
            base_name = os.path.splitext(file_name)[0]
            output_filename = f"{base_name}.json"
            output_path = os.path.join(result_dir, output_filename)
            
            save_to_json(data, output_path)
            
            print(f"✓ 解析成功: {file_name}")
            success_count += 1
            
        except Exception as e:
            print(f"✗ 解析失败: {file_name} - {e}")
    
    print(f"\n批量处理完成! 成功处理 {success_count}/{len(files_to_process)} 个文件")


def single_file_mode():
    """单文件处理模式主函数"""
    docs_dir = "docs"
    while True:
        print("\n文档解析器 - 单文件模式")
        print("=" * 50)
        
        # 获取文档文件
        all_files, _, _ = get_document_files(docs_dir)
        if not all_files:
            print("错误: docs目录中没有找到支持的文档文件")
            return
            
        # 显示文档列表
        print(f"\n找到 {len(all_files)} 个文档文件:")
        for i, file in enumerate(all_files, 1):
            print(f"  {i}. {file}")
            
        # 获取用户选择
        print("\n输入 'q' 退出单文件模式")
        choice = input("请选择要解析的文档编号: ").strip()
        
        if choice.lower() == 'q':
            print("退出单文件模式")
            break
            
        try:
            choice_idx = int(choice) - 1
            if choice_idx < 0 or choice_idx >= len(all_files):
                print("无效的选择")
                continue
                
            selected_file = all_files[choice_idx]
            file_path = os.path.join(docs_dir, selected_file)
            
            # 根据文件后缀调用相应的parser
            file_ext = os.path.splitext(selected_file)[1].lower()
            print(f"\n正在处理文档: {selected_file}")
            
            data = parse_document(file_path)
            
            # 保存结果
            result_dir = "result"
            os.makedirs(result_dir, exist_ok=True)
            base_name = os.path.splitext(selected_file)[0]
            output_filename = f"{base_name}.json"
            output_path = os.path.join(result_dir, output_filename)
            
            save_to_json(data, output_path)
            
            print("\n解析完成!")
            print(f"共找到 {len(data['sections'])} 个主要章节")
            print(f"结果已保存到: {output_path}")
            
        except ValueError:
            print("请输入有效的数字或 'q' 退出")
        except Exception as e:
            print(f"处理文档时出错: {e}")


def main():
    """主函数"""
    print("欢迎使用文档解析器")
    print("=" * 50)
    
    while True:
        print("\n请选择模式:")
        print("1. 单文件处理模式")
        print("2. 批量文件处理模式")
        print("3. 退出")
        
        try:
            choice = input("请选择 (1-3): ").strip()
            
            if choice == "1":
                single_file_mode()
            elif choice == "2":
                batch_mode_main()
            elif choice == "3":
                print("感谢使用，再见!")
                break
            else:
                print("无效选择，请重新输入")
                
        except KeyboardInterrupt:
            print("\n\n用户取消操作，感谢使用!")
            break
        except Exception as e:
            print(f"发生错误: {e}")
            continue


if __name__ == '__main__':
    main()